
import os
import logging

from migen import *

from litex.gen import *

from litex import get_data_mod

from litex.soc.interconnect import axi
from litex.soc.interconnect import wishbone

from litex.soc.integration.soc import SoCRegion

from litex.soc.cores.cpu import CPU, CPU_GCC_TRIPLE_RISCV64

# Xiangshan  ------------------------------------------------------------------------------------------

class Xiangshan(CPU):
    category             = "softcore"
    family               = "riscv"
    name                 = "xiangshan"
    human_name           = "XiangshanRV64[imac]"
    variants             = ["minimal", "medium", "default"]
    data_width           = 64
    endianness           = "little"
    gcc_triple           = CPU_GCC_TRIPLE_RISCV64
    linker_output_format = "elf64-littleriscv"
    nop                  = "nop"
    io_regions           = {0x4000_0000: 0x4000_0000} # Origin, Length.

    # Default parameters.
    cpu_num_cores = 1
    cpu_mem_width = 4

    # Command line configuration arguments.
    @staticmethod
    def args_fill(parser):
        cpu_group = parser.add_argument_group(title="CPU options")
        cpu_group.add_argument("--cpu-num-cores", default=1, help="Number of cores (1, 2, 4, 8).", type=int)
        cpu_group.add_argument("--cpu-mem-width", default=1, help="Width of memory port (1, 2, 4, 8).", type=int)

    @staticmethod
    def args_read(args):
        logger = logging.getLogger("XiangshanArgs")
        if int(args.cpu_num_cores) in [1, 2, 4, 8]:
            Xiangshan.cpu_num_cores = int(args.cpu_num_cores)
        else:
            logger.error("Invalid '--cpu-num-cores {}' (should be 1, 2, 4, or 8)".format(args.cpu_num_cores))
            quit()
        if int(args.cpu_mem_width) in [1, 2, 4, 8]:
            Xiangshan.cpu_mem_width = int(args.cpu_mem_width)
        else:
            logger.error("Invalid '--cpu-mem-width {}' (should be 1, 2, 4, or 8)".format(args.cpu_mem_width))
            quit()

    # Arch.
    @staticmethod
    def get_arch(variant):
        if variant in ["minimal", "medium", "default"]:
            return "rv64i2p0_mafdc"
        else:
            return "rv64i2p0_mac"

    # Memory Mapping.
    @property
    def mem_map(self):
        # Xiangshan reserves the first 256Mbytes for internal use, so we must change default mem_map.
        return {
            "rom"      : 0x1000_0000,
            "sram"     : 0x3000_0000,
            "clint"    : 0x3800_0000,
            "plic"     : 0x3c00_0000,
            "csr"      : 0x7000_0000,
            "main_ram" : 0x8000_0000,
        }

    # GCC Flags.
    @property
    def gcc_flags(self):
        flags =  "-mno-save-restore "
        flags += f"-march={self.get_arch(self.variant)} -mabi=lp64 "
        flags += "-D__xiangshan__ "
        flags += "-D__riscv_plic__ "
        flags += "-mcmodel=medany"
        return flags

    def __init__(self, platform, variant="minimal"):
        self.platform  = platform
        self.variant   = variant

        self.reset     = Signal()
        self.interrupt = Signal(64)

        mem_dw = 64 * Xiangshan.cpu_mem_width
        mmio_dw = 64

        self.mem_axi   =  mem_axi = axi.AXIInterface(data_width=mem_dw,  address_width=32, id_width=14)
        self.mmio_axi  = mmio_axi = axi.AXIInterface(data_width=mmio_dw, address_width=32, id_width=2)
        self.l2fb_axi  = l2fb_axi = axi.AXIInterface(data_width=mem_dw,  address_width=32, id_width=14)

        self.mmio_wb   = mmio_wb = wishbone.Interface(data_width=mmio_dw, adr_width=32-log2_int(mmio_dw//8), addressing="word")
        self.l2fb_wb   = l2fb_wb = wishbone.Interface(data_width=mem_dw,  adr_width=32-log2_int(mem_dw//8), addressing="word")

        self.memory_buses = [mem_axi] # Memory buses (Connected directly to LiteDRAM).
        self.periph_buses = [mmio_wb] # Peripheral buses (Connected to main SoC's bus).
        self.dma_bus      =  l2fb_wb  # DMA bus (Arbitrated and connected to SoC's bus).

        # # #
        self.peri_aw_addr = Signal(31)
        self.peri_ar_addr = Signal(31)
        self.comb += mmio_axi.aw.addr[0:31].eq(self.peri_aw_addr)
        self.comb += mmio_axi.ar.addr[0:31].eq(self.peri_ar_addr)
        self.comb += mmio_axi.aw.addr[31:].eq(Constant(0,1))
        self.comb += mmio_axi.ar.addr[31:].eq(Constant(0,1))

        self.cpu_params = dict(
            # Clk / Rst.
            i_io_clock = ClockSignal("sys"),
            i_io_reset = ResetSignal("sys") | self.reset,

            i_io_sram_config      = Constant(0,16),
            i_io_pll0_lock        = 0,
            o_io_pll0_ctrl_0      = Open(32),
            o_io_pll0_ctrl_1      = Open(32),
            o_io_pll0_ctrl_2      = Open(32),
            o_io_pll0_ctrl_3      = Open(32),
            o_io_pll0_ctrl_4      = Open(32),
            o_io_pll0_ctrl_5      = Open(32),

            # System Jtag Debug (ignored).
            o_io_debug_reset                      = Open(),
            i_io_systemjtag_reset                 = ResetSignal() | self.reset,
            i_io_systemjtag_mfr_id                = Constant(0,11),
            i_io_systemjtag_part_number           = Constant(0,16),
            i_io_systemjtag_version               = Constant(0,4),
            i_io_systemjtag_jtag_TCK              = 0,
            i_io_systemjtag_jtag_TMS              = 0,
            i_io_systemjtag_jtag_TDI              = 0,
            o_io_systemjtag_jtag_TDO_data         = Open(),
            o_io_systemjtag_jtag_TDO_driven       = Open(),

            i_io_cacheable_check_req_0_valid      = 0,
            i_io_cacheable_check_req_0_bits_addr  = Constant(0,36),
            i_io_cacheable_check_req_0_bits_size  = Constant(0,2),
            i_io_cacheable_check_req_0_bits_cmd   = Constant(0,3),
            i_io_cacheable_check_req_1_valid      = 0,
            i_io_cacheable_check_req_1_bits_addr  = Constant(0,36),
            i_io_cacheable_check_req_1_bits_size  = Constant(0,2),
            i_io_cacheable_check_req_1_bits_cmd   = Constant(0,3),

            o_io_cacheable_check_resp_0_ld      = Open(),
            o_io_cacheable_check_resp_0_st      = Open(),
            o_io_cacheable_check_resp_0_instr   = Open(),
            o_io_cacheable_check_resp_0_mmio    = Open(),
            o_io_cacheable_check_resp_1_ld      = Open(),
            o_io_cacheable_check_resp_1_st      = Open(),
            o_io_cacheable_check_resp_1_instr   = Open(),
            o_io_cacheable_check_resp_1_mmio    = Open(),
            o_io_riscv_halt_0                   = Open(),

            # IRQ.
            i_io_extIntrs = self.interrupt,

            # AXI Memory (L1-cached).
            i_memory_0_awready = mem_axi.aw.ready,
            o_memory_0_awvalid = mem_axi.aw.valid,
            o_memory_0_awid    = mem_axi.aw.id,
            o_memory_0_awaddr  = Cat(mem_axi.aw.addr,Signal(4)),
            o_memory_0_awlen   = mem_axi.aw.len,
            o_memory_0_awsize  = mem_axi.aw.size,
            o_memory_0_awburst = mem_axi.aw.burst,
            o_memory_0_awlock  = mem_axi.aw.lock,
            o_memory_0_awcache = mem_axi.aw.cache,
            o_memory_0_awprot  = mem_axi.aw.prot,
            o_memory_0_awqos   = mem_axi.aw.qos,

            i_memory_0_wready  = mem_axi.w.ready,
            o_memory_0_wvalid  = mem_axi.w.valid,
            o_memory_0_wdata   = mem_axi.w.data,
            o_memory_0_wstrb   = mem_axi.w.strb,
            o_memory_0_wlast   = mem_axi.w.last,

            o_memory_0_bready  = mem_axi.b.ready,
            i_memory_0_bvalid  = mem_axi.b.valid,
            i_memory_0_bid     = mem_axi.b.id,
            i_memory_0_bresp   = mem_axi.b.resp,

            i_memory_0_arready = mem_axi.ar.ready,
            o_memory_0_arvalid = mem_axi.ar.valid,
            o_memory_0_arid    = mem_axi.ar.id,
            o_memory_0_araddr  = Cat(mem_axi.ar.addr,Signal(4)),
            o_memory_0_arlen   = mem_axi.ar.len,
            o_memory_0_arsize  = mem_axi.ar.size,
            o_memory_0_arburst = mem_axi.ar.burst,
            o_memory_0_arlock  = mem_axi.ar.lock,
            o_memory_0_arcache = mem_axi.ar.cache,
            o_memory_0_arprot  = mem_axi.ar.prot,
            o_memory_0_arqos   = mem_axi.ar.qos,

            o_memory_0_rready  = mem_axi.r.ready,
            i_memory_0_rvalid  = mem_axi.r.valid,
            i_memory_0_rid     = mem_axi.r.id,
            i_memory_0_rdata   = mem_axi.r.data,
            i_memory_0_rresp   = mem_axi.r.resp,
            i_memory_0_rlast   = mem_axi.r.last,

            # AXI MMIO (not cached).
            i_peripheral_0_awready = mmio_axi.aw.ready,
            o_peripheral_0_awvalid = mmio_axi.aw.valid,
            o_peripheral_0_awid    = mmio_axi.aw.id,
            o_peripheral_0_awaddr  = self.peri_aw_addr,
            #o_peripheral_0_awaddr  = mmio_axi.aw.addr,
            o_peripheral_0_awlen   = mmio_axi.aw.len,
            o_peripheral_0_awsize  = mmio_axi.aw.size,
            o_peripheral_0_awburst = mmio_axi.aw.burst,
            o_peripheral_0_awlock  = mmio_axi.aw.lock,
            o_peripheral_0_awcache = mmio_axi.aw.cache,
            o_peripheral_0_awprot  = mmio_axi.aw.prot,
            o_peripheral_0_awqos   = mmio_axi.aw.qos,

            i_peripheral_0_wready  = mmio_axi.w.ready,
            o_peripheral_0_wvalid  = mmio_axi.w.valid,
            o_peripheral_0_wdata   = mmio_axi.w.data,
            o_peripheral_0_wstrb   = mmio_axi.w.strb,
            o_peripheral_0_wlast   = mmio_axi.w.last,

            o_peripheral_0_bready  = mmio_axi.b.ready,
            i_peripheral_0_bvalid  = mmio_axi.b.valid,
            i_peripheral_0_bid     = mmio_axi.b.id,
            i_peripheral_0_bresp   = mmio_axi.b.resp,

            i_peripheral_0_arready = mmio_axi.ar.ready,
            o_peripheral_0_arvalid = mmio_axi.ar.valid,
            o_peripheral_0_arid    = mmio_axi.ar.id,
            o_peripheral_0_araddr  = self.peri_ar_addr,
            #o_peripheral_0_araddr  = mmio_axi.ar.addr,
            o_peripheral_0_arlen   = mmio_axi.ar.len,
            o_peripheral_0_arsize  = mmio_axi.ar.size,
            o_peripheral_0_arburst = mmio_axi.ar.burst,
            o_peripheral_0_arlock  = mmio_axi.ar.lock,
            o_peripheral_0_arcache = mmio_axi.ar.cache,
            o_peripheral_0_arprot  = mmio_axi.ar.prot,
            o_peripheral_0_arqos   = mmio_axi.ar.qos,

            o_peripheral_0_rready  = mmio_axi.r.ready,
            i_peripheral_0_rvalid  = mmio_axi.r.valid,
            i_peripheral_0_rid     = mmio_axi.r.id,
            i_peripheral_0_rdata   = mmio_axi.r.data,
            i_peripheral_0_rresp   = mmio_axi.r.resp,
            i_peripheral_0_rlast   = mmio_axi.r.last,

            # AXI L2FB (Slave, for e.g., DMA).
            o_dma_0_awready = l2fb_axi.aw.ready,
            i_dma_0_awvalid = l2fb_axi.aw.valid,
            i_dma_0_awid    = l2fb_axi.aw.id,
            i_dma_0_awaddr  = Cat(l2fb_axi.aw.addr,Constant(0,4)),
            i_dma_0_awlen   = l2fb_axi.aw.len,
            i_dma_0_awsize  = l2fb_axi.aw.size,
            i_dma_0_awburst = l2fb_axi.aw.burst,
            i_dma_0_awlock  = l2fb_axi.aw.lock,
            i_dma_0_awcache = l2fb_axi.aw.cache,
            i_dma_0_awprot  = l2fb_axi.aw.prot,
            i_dma_0_awqos   = l2fb_axi.aw.qos,

            o_dma_0_wready  = l2fb_axi.w.ready,
            i_dma_0_wvalid  = l2fb_axi.w.valid,
            i_dma_0_wdata   = l2fb_axi.w.data,
            i_dma_0_wstrb   = l2fb_axi.w.strb,
            i_dma_0_wlast   = l2fb_axi.w.last,

            i_dma_0_bready  = l2fb_axi.b.ready,
            o_dma_0_bvalid  = l2fb_axi.b.valid,
            o_dma_0_bid     = l2fb_axi.b.id,
            o_dma_0_bresp   = l2fb_axi.b.resp,

            o_dma_0_arready = l2fb_axi.ar.ready,
            i_dma_0_arvalid = l2fb_axi.ar.valid,
            i_dma_0_arid    = l2fb_axi.ar.id,
            i_dma_0_araddr  = Cat(l2fb_axi.ar.addr,Constant(0,4)),
            i_dma_0_arlen   = l2fb_axi.ar.len,
            i_dma_0_arsize  = l2fb_axi.ar.size,
            i_dma_0_arburst = l2fb_axi.ar.burst,
            i_dma_0_arlock  = l2fb_axi.ar.lock,
            i_dma_0_arcache = l2fb_axi.ar.cache,
            i_dma_0_arprot  = l2fb_axi.ar.prot,
            i_dma_0_arqos   = l2fb_axi.ar.qos,

            i_dma_0_rready  = l2fb_axi.r.ready,
            o_dma_0_rvalid  = l2fb_axi.r.valid,
            o_dma_0_rid     = l2fb_axi.r.id,
            o_dma_0_rdata   = l2fb_axi.r.data,
            o_dma_0_rresp   = l2fb_axi.r.resp,
            o_dma_0_rlast   = l2fb_axi.r.last,
        )
        # additional per-core debug signals:
        #self.cpu_params.update({'i_resetctrl_hartIsInReset_%s'%i : Open() for i in range(Xiangshan.cpu_num_cores)})

        # Adapt AXI interfaces to Wishbone.
        mmio_a2w = axi.AXI2Wishbone(mmio_axi, mmio_wb, base_address=0)
        self.submodules += mmio_a2w

        l2fb_a2w = axi.Wishbone2AXI(l2fb_wb, l2fb_axi, base_address=0)
        self.submodules += l2fb_a2w

        # Add Verilog sources.
        self.add_sources(platform, variant)

    def set_reset_address(self, reset_address):
        self.reset_address = reset_address
        assert reset_address == 0x1000_0000, "cpu_reset_addr hardcoded in during elaboration!"

    @staticmethod
    def add_sources(platform, variant):
        pfx = "xiangshan.system.LitexConfig"
        fname = f"{pfx}_{variant}_{Xiangshan.cpu_num_cores}_{Xiangshan.cpu_mem_width}"
        vdir = get_data_mod("cpu", "xiangshan").data_location
        platform.add_sources(
            os.path.join(vdir, "vsrc"),
            "define.v",
            #"AsyncResetReg.v",
            #"EICG_wrapper.v",
        )
        platform.add_sources(
            os.path.join(vdir, "generated-src"),
            f"{fname}.v",
            f"{fname}.behav_srams.v",
        )

    def add_soc_components(self, soc):
        # Add OpenSBI/PLIC/CLINT regions.
        soc.bus.add_region("opensbi", SoCRegion(origin=self.mem_map["main_ram"] + 0x0000_0000, size=0x20_0000, cached=True, linker=True))
        soc.bus.add_region("plic",    SoCRegion(origin=soc.mem_map.get("plic"),               size=0x40_0000, cached=True,  linker=True))
        soc.bus.add_region("clint",   SoCRegion(origin=soc.mem_map.get("clint"),               size= 0x1_0000, cached=True,  linker=True))

        # Define number of CPUs
        soc.add_config("CPU_COUNT", Xiangshan.cpu_num_cores)
        soc.add_config("CPU_ISA",   self.get_arch(self.variant))
        soc.add_config("CPU_MMU",   "sv39")

        # Constants for Cache so we can add them in the DTS.
        soc.add_config("CPU_DCACHE_SIZE",      65536) # CHECKME: correct/hardwired?
        soc.add_config("CPU_DCACHE_WAYS",          4) # CHECKME: correct/hardwired?
        soc.add_config("CPU_DCACHE_BLOCK_SIZE",   64) # CHECKME: correct/hardwired?

        soc.add_config("CPU_ICACHE_SIZE",      65536) # CHECKME: correct/hardwired?
        soc.add_config("CPU_ICACHE_WAYS",          4) # CHECKME: correct/hardwired?
        soc.add_config("CPU_ICACHE_BLOCK_SIZE",   64) # CHECKME: correct/hardwired?

        # Constants for TLB so we can add them in the DTS.
        soc.add_config("CPU_DTLB_SIZE", 64) # CHECKME: correct/hardwired?
        soc.add_config("CPU_DTLB_WAYS",  1) # CHECKME: correct/hardwired?

        soc.add_config("CPU_ITLB_SIZE", 32) # CHECKME: correct/hardwired?
        soc.add_config("CPU_ITLB_WAYS",  1) # CHECKME: correct/hardwired?

    def do_finalize(self):
        assert hasattr(self, "reset_address")
        self.specials += Instance("XSTop", **self.cpu_params)
